# --------------------------------------------------------------------------
# Version 2 of simple blackboard application that generates integer values,
# computes their squares, and prints the squares.
# Will evolve this to exercise more of the new blackboard functionality
# --------------------------------------------------------------------------

import blackboard as blackboard
from BB_Observations import BB_Observations;
import scipy.stats as stats;
import pgm, warnings;  # @UnusedImport
import numpy as np;
import hmm_ext;
import pandas as pd;
import pymc as pm;
import pylab as pl;
# from b3.blackboard import *

STOP_VALUE = 5
"""
Specifies the last integer generated by the generate_integers KS
"""

VERBOSE = True
"""
When True, prints KS.execution_function activity for debugging
"""

# BB_Observations Instance
FILE_PATH_DB = '../../../data/PoL.s3db'; #'./PoL.s3db';
bb_Detections = BB_Observations(FILE_PATH_DB);


# --------------------------------------------------------------------------
# A simple "blackboard" application that reads raw data from sensors and
# generates observation datasets for inividual and group tracks. 
# --------------------------------------------------------------------------

# --------------------------------------------------------------------------
# Events

def store_table(dataFrame, sTblName, **kwargs):
    """
    Stores a table 'dataFrame' with name 'sTblName' into the DB.
    Signals a database transaction event in blackboard.
    It should be done in the context of a blackboard, which could either be
    passed as a keyword argument: blackboard=<bb-instance> (higher precendence)
    OR when global CURRENT_BLACKBOARD reference the current blackboard instance. 
    'sTblName' must be a string.
    """
#     dataFrame = unit_db_instance.value;
    global VERBOSE
    if VERBOSE:
        print ">>> Storing Table {0}: {1}".format( sTblName, type(dataFrame) );
    CURRENT_BLACKBOARD = blackboard.get_current_blackboard();
    if CURRENT_BLACKBOARD and 'blackboard' not in kwargs:
        bb_Detections.storeTable(dataFrame, sTblName);
        CURRENT_BLACKBOARD.signal_event(blackboard.DB_Event, #unit_instance=unit_db_instance, 
                                        sDB_Event_Type="Create",
                                        sTblName=sTblName); #"Update_Table_"+sTblName, sTblName);
    else:
        bb_Detections.storeTable(dataFrame, sTblName);


# --------------------------------------------------------------------------
# Units

class Raw_Data(blackboard.Standard_Unit_Instance):
    """
    A blackboard object containing the Raw Data from Sensors.
    """
    type_name = 'Raw_Data'
    def __init__(self, **kwargs):
        self.value = None
        if 'value' in kwargs:
            self.value = kwargs['value']
    def __repr__(self):
        return 'Raw_Data({0})'.format(type(self.value))
    
class IndObs_Data(blackboard.Standard_Unit_Instance):
    """
    A blackboard object containing the data about Individual Observation 
    Attributes generated from Generate_IndObs_KS.
    """
    type_name = 'IndObs_Data'
    def __init__(self, **kwargs):
        self.value = None
        if 'value' in kwargs:
            self.value = kwargs['value']
    def __repr__(self):
        return 'IndObs_Data({0})'.format(type(self.value))
    
class GrpObs_Data(blackboard.Standard_Unit_Instance):
    """
    A blackboard object containing the data about Groups Observation 
    Attributes generated from Generate_GrpObs_KS.
    """
    type_name = 'GrpObs_Data'
    def __init__(self, **kwargs):
        self.value = None
        if 'value' in kwargs:
            self.value = kwargs['value']
    def __repr__(self):
        return 'GrpObs_Data({0})'.format(type(self.value))
    
class IndGrp_Data(blackboard.Standard_Unit_Instance):
    """
    A blackboard object containing the data about Individual & Groups Observation 
    Attributes generated from Generate_GrpObs_KS.
    """
    type_name = 'IndGrp_Data'
    def __init__(self, **kwargs):
        self.value = None
        if 'value' in kwargs:
            self.value = kwargs['value']
    def __repr__(self):
        return 'IndGrp_Data({0})'.format(type(self.value))
    
class PairObs_Data(blackboard.Standard_Unit_Instance):
    """
    A blackboard object containing the data about Pair-wise Observation 
    Attributes generated from Generate_GrpObs_KS.
    """
    type_name = 'PairObs_Data'
    def __init__(self, **kwargs):
        self.value = None
        if 'value' in kwargs:
            self.value = kwargs['value']
    def __repr__(self):
        return 'PairObs_Data({0})'.format(type(self.value))
    
class Gait_IndLM_Data(blackboard.Standard_Unit_Instance):
    """
    A blackboard object containing the Local Model about Individual Gait
    like Run, Walk, Stand with Naive Bayes considering Observation Speed.
    """
    type_name = 'Gait_IndLM_Data'
    def __init__(self, **kwargs):
        self.value = None
        if 'value' in kwargs:
            self.value = kwargs['value']
    def __repr__(self):
        return 'Gait_IndLM_Data({0})'.format(type(self.value))
    
class Gait_GrpLM_Data(blackboard.Standard_Unit_Instance):
    """
    A blackboard object containing the Local Model about Group's Gait
    like Run, Walk, Stand with Naive Bayes considering Observation Speed.
    """
    type_name = 'Gait_GrpLM_Data'
    def __init__(self, **kwargs):
        self.value = None
        if 'value' in kwargs:
            self.value = kwargs['value']
    def __repr__(self):
        return 'Gait_GrpLM_Data({0})'.format(type(self.value))


# --------------------------------------------------------------------------
# KSSs

class Startup_KS(blackboard.Knowledge_Source):
    """
    Startup KS
    Triggered by Control_Shell_Started_Event
    """
    def __init__(self, **kwargs):
        blackboard.Knowledge_Source.__init__(self,
                                             trigger_events=( blackboard.Control_Shell_Started_Event, ),
                                             execution_function=self.startup_ks_function,
                                             **kwargs)
        
    def startup_ks_function(self, ksa):
        """
        KSA Startup_KS execution function:
        Creates a single Integer_Unit on the BB
        """
        event = ksa.ksa_execution_context[0]
    
        global VERBOSE
        if VERBOSE:
            print ">>> Executing startup_ks_function({0})".format(ksa)
            print ">>>     ksa.execution_cycle: {0}".format(ksa.execution_cycle)
            print ">>>     ksa.rating: {0}".format(ksa.rating)
            print ">>>     ksa.ksa_execution_context: {0}".format(ksa.ksa_execution_context)
            print ">>>     event: {0}".format(event)
            print ">>>     Calling make_unit_instance(Integer_Unit, value=1)"
            
        #Get Raw Data
        nInitFrm, nEndFrm = 0, 25;
        warnings.filterwarnings("ignore");
        dfRawData = bb_Detections.getRawData(FILE_PATH_DB, nInitFrm, nEndFrm);
        blackboard.make_unit_instance(Raw_Data, value=dfRawData);
#         make_bb_object("Raw_Data", Raw_Data(dfRawData));    
#         blackboard.make_unit_instance(Integer_Unit, value=1)



class Generate_IndObs_KS(blackboard.Knowledge_Source):
    """ 
    KS that generates a dataset IndObs_Data about Individual Observation 
    Attributes obtained from using the dataset of Raw_Data from Sensors.
    """
    def __init__(self, **kwargs):
        blackboard.Knowledge_Source.__init__(self,
                                             trigger_events=( (blackboard.Instance_Created_Event, Raw_Data), ),
                                             execution_function=self.generate_indobs_ks_function,
                                             **kwargs);
    def generate_indobs_ks_function(self, ksa):
        event = ksa.ksa_execution_context[0];
        dfRawData = event.payload["unit_instance"].value;
        dfIndAttrs = bb_Detections.getIndAttrs(dfRawData.copy());
        unit_IndObs_Data = blackboard.make_unit_instance(IndObs_Data, value=dfIndAttrs)  # @UnusedVariable
        store_table(dfIndAttrs, "IndAttrs");

        
class Generate_GrpObs_KS(blackboard.Knowledge_Source):
    """ 
    KS that generates the datasets GrpObs_Data (about Groups Observation 
    Attributes) and Grp_Data (about Identifying Membership of each individual
    track to their respective group). Both obtained from using the 
    dataset of IndObs_Data generated from Generate_IndObs_KS.
    Then it generates the dataset IndGrp_Data (info about Group & Individual
    Observation's average positions).
    """
    def __init__(self, **kwargs):
        blackboard.Knowledge_Source.__init__(self,
                                             trigger_events=( (blackboard.Instance_Created_Event, IndObs_Data), ),
                                             execution_function=self.generate_grpobs_ks_function,
                                             **kwargs);
    def generate_grpobs_ks_function(self, ksa):
        event = ksa.ksa_execution_context[0];
        dfIndAttrs = event.payload["unit_instance"].value;
        CURRENT_BLACKBOARD = blackboard.get_current_blackboard();  # @UnusedVariable
#         obj = CURRENT_BLACKBOARD.get_instance_bookkeeping_state(instance_class)
        [dfGroups, dfGrp_Attrs] = bb_Detections.getGrpAttrs(dfIndAttrs);  # @UnusedVariable
        unit_GrpObs_Data = blackboard.make_unit_instance(GrpObs_Data, value=[dfGroups, dfGrp_Attrs]);  # @UnusedVariable
#         unit_GrpObs_Data = blackboard.make_unit_instance(GrpObs_Data, value=dfGrp_Attrs);  # @UnusedVariable
        dfIndGrp_Attrs = bb_Detections.getIndGrp_Attrs(dfIndAttrs, dfGrp_Attrs);
        blackboard.make_unit_instance(IndGrp_Data, value=dfIndGrp_Attrs);
        store_table(dfGroups, "IndAttrs");
        store_table(dfGrp_Attrs, "GrpAttrs");
 
         
class Generate_PairObs_KS(blackboard.Knowledge_Source):
    """ 
    KS that generates the dataset PairObs_Data about Pair-Wise Observation 
    Attributes for Groups and Individual tracks) obtained from using the 
    dataset of IndGrp_Data generated from Generate_GrpObs_KS.
    """
    def __init__(self, **kwargs):
        blackboard.Knowledge_Source.__init__(self,
                                             trigger_events=( (blackboard.Instance_Created_Event, IndGrp_Data), ),
                                             execution_function=self.generate_pairobs_ks_function,
                                             **kwargs);
    def generate_pairobs_ks_function(self, ksa):
        event = ksa.ksa_execution_context[0];
        dfIndGrp_Attrs = event.payload["unit_instance"].value;
        dfPairAttrs = bb_Detections.getPairAttrs(dfIndGrp_Attrs);
        unit_PairObs_Data = blackboard.make_unit_instance(PairObs_Data, value=dfPairAttrs);  # @UnusedVariable
        store_table(dfPairAttrs, "PairAttrs");


class Gait_IndLM_KS(blackboard.Knowledge_Source):
    def __init__(self, **kwargs):
        self.nThreshold = 0.1;
        blackboard.Knowledge_Source.__init__(self,
                                             trigger_events=( (blackboard.DB_Event, IndObs_Data), ),
                                             execution_function=self.gait_indlm_ks_function,
                                             **kwargs);        
    def gait_indlm_ks_function(self, ksa):
        global VERBOSE;
        event = ksa.ksa_execution_context[0];
        dfIndAttrs = event.payload["unit_instance"].value;
        if VERBOSE:
            print "\tGait_IndLM_KS is Active."
        for nTrackId, dfTrack in dfIndAttrs.groupby(['TrackId']):  # @UnusedVariable
            avgSpeed = dfTrack["Speed"].sum()/float(dfTrack["Speed"].count());
            nProb_Run = self.getRun(avgSpeed);
            nProb_Stand = self.getStand(avgSpeed);
            nProb_Walk = self.getWalk(avgSpeed);
            aProbs = {"Stand":nProb_Stand, "Walk":nProb_Walk, "Run":nProb_Run};
            for sGait in aProbs:
                if aProbs[sGait]>self.nThreshold:
                    data = self.getData(dfTrack, nTrackId, avgSpeed, aProbs[sGait], sGait);
                    sAvgSpeed = "%.2f" % avgSpeed;
                    sProb = "%.2f" % aProbs[sGait];
                    if VERBOSE:
                        print "\t** Creating New Data for ", [int(nTrackId), sGait, sProb, self.nThreshold, sAvgSpeed];
                    blackboard.make_unit_instance(Gait_IndLM_Data, value=data);
          
    def getData(self, dfTrack, nTrackId, avgSpeed, nProb_Gait, sGait):
#         g = pgm.Graph();
#         cpt1 = [.5, .5];
#         cpt2 = {"['False']": [.5, .5],"['True']": [nProb_Gait, 1-nProb_Gait]};
#         g.addnode(pgm.Node(sGait, ["False", "True"], [None], cpt1));
#         g.addnode(pgm.Node("Speed=%.2f"%avgSpeed, ["False", "True"], [g.node[sGait]], cpt2));
#         g.setup();
        G_obs = [1.];
        N = len(G_obs);
        gait = pm.Categorical(sGait, [0.5, 0.5], value=pl.ones(N));
        p_speed = pm.Lambda('p_'+sGait, lambda gait=gait: pl.where(gait, nProb_Gait, [0.5, 0.5]));
        speed = pm.Categorical("Speed=%.2f"%avgSpeed, p_speed, value=G_obs, observed=True);
        model = pm.Model([gait, speed]);
        g = pm.graph.graph(model);
        g.write_pdf("./Models/Graph2_"+str(int(nTrackId))+"_"+str(sGait)+".pdf");
#         g.write2pdf("./Models/Graph_"+str(int(nTrackId))+"_"+str(sGait)+".pdf");
        data = {"TrackId":nTrackId, "Type":sGait, "Belief":nProb_Gait, 
                "Obs":["Speed"], "Obs_Vals":[avgSpeed], "MEs":["Walk","Stand"], 
                "Graph":g};
        return data;
          
    def getRun(self, avgSpeed):
        nRun = abs(64 - avgSpeed*10);
        nProb_Run = (1 - stats.beta.cdf(nRun, 3,4));
        if avgSpeed > 64: nProb_Run = 1;
        return nProb_Run;
      
    def getStand(self, avgSpeed):
        nStand = avgSpeed*10;
        nProb_Stand = (1 - stats.beta.cdf(nStand, 3,4));
        nProb_Stand = (0.6 + nProb_Stand*0.4);
        return nProb_Stand;
      
    def getWalk(self, avgSpeed):
        nWalk = abs(32 - avgSpeed*10);
        nProb_Walk = (1 - stats.beta.cdf(nWalk, 3,4));
        nProb_Walk = (0.6 + nProb_Walk*0.4);
        return nProb_Walk;
    

class Gait_GrpLM_KS(blackboard.Knowledge_Source):
    def __init__(self, **kwargs):
        self.nThreshold = 0.1;
        blackboard.Knowledge_Source.__init__(self,
                                             trigger_events=( (blackboard.Instance_Created_Event, GrpObs_Data), ),
                                             execution_function=self.gait_grplm_ks_function,
                                             **kwargs);        
    def gait_grplm_ks_function(self, ksa):
        global VERBOSE;
        event = ksa.ksa_execution_context[0];
        [dfGroups, dfGrp_Attrs] = event.payload["unit_instance"].value;
        if VERBOSE:
            print "\tGait_GrpLM_KS is Active."
        dfGrps = dfGrp_Attrs[pd.notnull(dfGrp_Attrs["EndFrm"])];
        dfGrps = dfGrps[["GrpId", "InitFrm", "EndFrm"]];
        aSegments = [];
        aInitFrms = [];
        for idx in dfGrps.index: #For each Grp Segment
            nGrpId = int(dfGrps["GrpId"][idx]);
            nInitFrm, nEndFrm = int(dfGrps["InitFrm"][idx]), int(dfGrps["EndFrm"][idx]);
            aInitFrms.append(nInitFrm);
            dfGrp_Segment = dfGroups[(dfGroups["GrpId"]==nGrpId) & (dfGroups["Frame"]>=nInitFrm) & 
                           (dfGroups["Frame"]<=nEndFrm)][["TrackId", "Frame", "Speed"]];
            aGrpSgm = [];
            for nFrm, dfFrm in dfGrp_Segment.groupby(["Frame"]):  # @UnusedVariable
                aFrmData = dfFrm["Speed"].tolist();
                aGrpSgm.append(aFrmData);
            aSegments.append(aGrpSgm);
#         print "HMM-1 Frames:", len(aSegments[0]), "InitFrame:", aInitFrms[0];
        data = self.getHMM_Gait(aSegments[0], aInitFrms[0]);
        if VERBOSE:
            print "\t** Creating New Data for ", data;
        blackboard.make_unit_instance(Gait_GrpLM_Data, value=data);
         
    def getHMM_Gait(self, obs_seq, nInitFrm):
        states = ["Stand", "Walk", "Run"];
        n_states = len(states);        
        start_probability = np.array([0.3, 0.5, 0.2]);
        transition_probability = np.array([
          [0.6, 0.3, 0.1],
          [0.3, 0.5, 0.2],
          [0.2, 0.3, 0.5]
        ]);        
        model = hmm_ext.NB_GammaHMM(n_components=n_states);
        model._set_startprob(start_probability);
        model._set_transmat(transition_probability);
        model._set_shapes(np.array([1,5,30]));
        model._set_locs(np.array([0,0,0]));
        model._set_scales(np.array([1,2,1]));        
        #Predict a sequence of hidden states based on visible states
        obs_seq = np.array(obs_seq)*3;
        logprob, hidden_states = model.decode(obs_seq, algorithm="viterbi"); # @UnusedVariable
        aHidden_States = map(lambda x: states[x], hidden_states);
#         print "Observations:", obs_seq;
#         print "Hidden States:", ", ".join(aHidden_States);
#         print "Posterior:", np.exp(logprob);
        #Get Sequence of Group Activity Segments 
        aSgmActs, aActData = [], [];
        sInitHiddenSt = aHidden_States[0];
        nInitSgm = nInitFrm;
        for i,sHiddenSt in enumerate(aHidden_States):
            if i==0: continue;
            if sHiddenSt!=sInitHiddenSt or i==len(aHidden_States)-1:
                aActData = (sInitHiddenSt, nInitSgm, (nInitFrm+i)-1);
                aSgmActs.append(aActData);
                sInitHiddenSt = sHiddenSt;
                nInitSgm = nInitFrm+i;
        return aSgmActs;
    
    
         
         
# # class Gait_GrpLM_KS(blackboard.Knowledge_Source):
# #     def __init__(self, **kwargs):
# #         self.nThreshold = 0.1;
# #         blackboard.Knowledge_Source.__init__(self,
# #                                              trigger_events=( (blackboard.DB_Event, GrpObs_Data), ),
# #                                              execution_function=self.gait_grplm_ks_function,
# #                                              **kwargs);        
# #     def gait_grplm_ks_function(self, ksa):
# #         global VERBOSE;
# #         event = ksa.ksa_execution_context[0];
# #         dfGrp_Attrs = event.payload["unit_instance"].value[1];
# #         if VERBOSE:
# #             print "\tGait_GrpLM_KS is Active."
# #         for nGrpId, dfGrp in dfGrp_Attrs.groupby(['GrpId']):  # @UnusedVariable
# #             avgSpeed = dfGrp["Speed"].sum()/float(dfGrp["Speed"].count());
# #             nProb_Run = self.getRun(avgSpeed);
# #             nProb_Stand = self.getStand(avgSpeed);
# #             nProb_Walk = self.getWalk(avgSpeed);
# #             aProbs = {"Stand":nProb_Stand, "Walk":nProb_Walk, "Run":nProb_Run};
# #             for sGait in aProbs:
# #                 if aProbs[sGait]>self.nThreshold:
# #                     data = self.getData(dfGrp, nGrpId, avgSpeed, aProbs[sGait], sGait);
# #                     sAvgSpeed = "%.2f" % avgSpeed;
# #                     sProb = "%.2f" % aProbs[sGait];
# #                     if VERBOSE:
# #                         print "\t** Creating New Data for ", [int(nGrpId), sGait, sProb, self.nThreshold, sAvgSpeed];
# #                     blackboard.make_unit_instance(Gait_GrpLM_Data, value=data);
# #          
# #     def getData(self, dfTrack, nGrpId, avgSpeed, nProb_Gait, sGait):
# # #         g = pgm.Graph();
# # #         cpt1 = [.5, .5];
# # #         cpt2 = {"['False']": [.5, .5],"['True']": [nProb_Gait, 1-nProb_Gait]};
# # #         g.addnode(pgm.Node(sGait, ["False", "True"], [None], cpt1));
# # #         g.addnode(pgm.Node("Speed=%.2f"%avgSpeed, ["False", "True"], [g.node[sGait]], cpt2));
# # #         g.setup();
# #         G_obs = [1.];
# #         N = len(G_obs);
# #         gait = pm.Bernoulli(sGait, .5, value=pl.ones(N));
# #         p_speed = pm.Lambda('p_'+sGait, lambda gait=gait: pl.where(gait, nProb_Gait, .5));
# #         speed = pm.Bernoulli("Speed=%.2f"%avgSpeed, p_speed, value=G_obs, observed=True);
# #         model = pm.Model([gait, speed]);
# #         g = pm.graph.graph(model);
# #         g.write_pdf("./Models/Graph2_"+str(int(nGrpId))+"_"+str(sGait)+".pdf");
# # #         g.write2pdf("./Models/Graph_Grp_"+str(int(nGrpId))+"_"+str(sGait)+".pdf");
# #         data = {"GrpId":nGrpId, "Type":sGait, "Belief":nProb_Gait, 
# #                 "Obs":["Speed"], "Obs_Vals":[avgSpeed], "MEs":["Walk","Stand"], 
# #                 "Graph":g};
# #         return data;
# #          
# #     def getRun(self, avgSpeed):
# #         nRun = abs(64 - avgSpeed*10);
# #         nProb_Run = (1 - stats.beta.cdf(nRun, 3,4));
# #         if avgSpeed > 64: nProb_Run = 1;
# #         return nProb_Run;
# #      
# #     def getStand(self, avgSpeed):
# #         nStand = avgSpeed*10;
# #         nProb_Stand = (1 - stats.beta.cdf(nStand, 3,4));
# #         nProb_Stand = (0.6 + nProb_Stand*0.4);
# #         return nProb_Stand;
# #      
# #     def getWalk(self, avgSpeed):
# #         nWalk = abs(32 - avgSpeed*10);
# #         nProb_Walk = (1 - stats.beta.cdf(nWalk, 3,4));
# #         nProb_Walk = (0.6 + nProb_Walk*0.4);
# #         return nProb_Walk;

# ---------------------------------------------------------------------------------------------------
# ---------------------------------------------------------------------------------------------------

def application_reset():
    # set up fresh environment
    blackboard.delete_workspace_image()
    bb = blackboard.create_blackboard_repository(make_current=True)

    # Enable event printing for _all_ events
    bb.enable_event_printing((blackboard.Standard_Event_Instance, True))

    # for turning on all Control_Shell_Events -- this is
    # covered by Standard_Event_Instance
    # bb.enable_event_printing((Control_Shell_Event, True))
    # The above line includes turning on event printing
    # for all KSA_Events.  As an example of disabling
    # subevents, the following turns off all KSA_Events
    # and their subevents, but does not affect other
    # Control_Shell_Events
    # bb.disable_event_printing((KSA_Event, True))

    # Define the KSs
    blackboard.define_knowledge_source(Startup_KS)
    blackboard.define_knowledge_source(Generate_IndObs_KS)
    blackboard.define_knowledge_source(Generate_GrpObs_KS)
    blackboard.define_knowledge_source(Generate_PairObs_KS)
    blackboard.define_knowledge_source(Gait_IndLM_KS)
    blackboard.define_knowledge_source(Gait_GrpLM_KS)


def run_application():
    application_reset()

    # Create the agenda control shell and start the control-loop
    acs = blackboard.Agenda_Control_Shell(activity_printing=True)

    print "\nBlackboard state before starting:"
    acs.blackboard.describe(verbose=True)
    acs.describe()

    result = acs.start()

    print "\nBlackboard state after stopping:"
    acs.blackboard.describe(verbose=True)
    acs.describe()

    return result



# --------------------------------------------------------------------------
# MAIN
# --------------------------------------------------------------------------

if __name__ == "__main__":
    run_application()